import mnist
import experiment_runner
import numpy as np # 导入 numpy 用于生成范围
# --- 1. 加载 MNIST 数据 ---
root_data,root_label,clients_data,clients_label,x_test,y_test = mnist.Load_MNIST(q = 0.1, nclient = 100, root_dataset_size = 100, root_dataset_bias = 0.1)

# --- 2. 定义实验参数 ---
smarter_attack_params = {
    'initial_benign_count_param': 50,
    'malicious_config_param': {'poor_benign': 20, 'good_benign': 50}, # 对应服务器方法中的malicious_config
    'benign_effectiveness_threshold_param': 0.5, # 示例阈值
    'cos_candidates_param': np.arange(0, 0.5, 0.05).tolist(), # 传递给服务器的cos候选列表，确保是可序列化的类型如果需要
    'expr_id_suffix': "smarter_attack_v1" # 给本次特定配置的实验一个标识
}
# --- 3. 运行基准实验 (先运行这些) ---
print("\n--- 运行基准: 无攻击 ---")
server_no_attack = experiment_runner.run_no_attacks(root_data,root_label,clients_data,clients_label, 100, x_test, y_test)

#print("\n--- 运行基准: 仅服务器训练 ---")
#server_only_server = experiment_runner.run_only_server(root_data,root_label, x_test, y_test)

# --- 4. 运行新的自适应攻击训练 ---
print(f"\n--- 开始训练更智能的自适应攻击模型 ({smarter_attack_params['expr_id_suffix']}) ---")

# 如果是从外部文件加载数据，确保 nclient 参数与 clients_data 的长度一致
# n_clients_from_data = len(clients_data) # 例如

trained_server_smarter = experiment_runner.run_smarter_adaptive_attack(
    root_data, root_label,
    clients_data, clients_label,
    nclient=len(clients_data), # 确保nclient与加载的客户端数据一致
    x_test=x_test, y_test=y_test,
    **smarter_attack_params # 使用字典解包传递参数
)
print(f"--- 完成训练 ({smarter_attack_params['expr_id_suffix']}) ---")

# --- 5. 最终评估阶段 ---
# ... (评估基准模型)
server_no_attack.evaluate(x_test, y_test, 'FLTrust in test data')
# 评估新的攻击模型
eval_label_smarter = f'[更智能攻击] FLTrust 测试集 ({smarter_attack_params["expr_id_suffix"]})'
trained_server_smarter.evaluate(x_test, y_test, eval_label_smarter)